from torch import nn


class UNetPlusPlusLoss(nn.Module):
    def __init__(self, ce_weight=0.5, dice_weight=0.5):
        super().__init__()
        self.ce_weight = ce_weight
        self.dice_weight = dice_weight

        # 如果输出已经经过sigmoid，应该使用BCE而不是BCEWithLogitsLoss
        self.ce = nn.BCELoss()  # 修改这里

    def dice_loss(self, pred, target):
        # 这里不需要再做sigmoid了，因为输入已经是sigmoid后的结果
        pred = pred.view(-1)
        target = target.view(-1)

        intersection = (pred * target).sum()
        union = pred.sum() + target.sum()

        dice = 1 - (2.0 * intersection + 1e-7) / (union + 1e-7)
        return dice

    def cul_loss(self, pred, target):
        # pred已经是sigmoid后的结果，直接用于计算loss
        ce_loss = self.ce(pred, target)
        dice_loss = self.dice_loss(pred, target)

        total_loss = self.ce_weight * ce_loss + self.dice_weight * dice_loss
        return total_loss
